Repository navigation
fix(spanner): record streaming operation completion when stream finishes - #18704
sinhasubham wants to merge 1 commit into
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request refactors the Spanner client's metrics tracking to ensure that streaming operation latencies are recorded immediately upon stream completion and exclude consumer processing delays. This is achieved by tracking stream completion in both synchronous and asynchronous snapshot iterators and recording metrics accordingly. However, the reviewer identified a critical issue in metrics_interceptor.py where the code attempts to check for a last attribute on the yielded stream items. Since the production PartialResultSet protobuf messages do not contain a last field, this check will always evaluate to False and act as dead code. It is recommended to remove these redundant checks.
fcc2c32 to
1b8a925
Compare
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request updates the Spanner client to record operation completion metrics immediately when the stream finishes or when a PartialResultSet with last=True is observed, rather than waiting for background stream draining. The feedback suggests simplifying the boolean checks in metrics_interceptor.py by removing the explicit 'is True' comparisons to align with PEP 8 guidelines.
1b8a925 to
31b7a08
Compare
…empt latency upon stream completion Wait until the streaming response completes (either via PartialResultSet.last=True or iterator exhaustion/closure) before recording operation completion and attempt completion metrics in _restart_on_unavailable and _StreamingResponseWrapper / _AsyncStreamingResponseWrapper. Fixes #18656
31b7a08 to
e37b728
Compare
| items_done = [it async for it in resumable_done] | ||
| self.assertEqual(items_done, [item_two]) | ||
| mock_tracer.record_operation_completion.assert_called_once() |
There was a problem hiding this comment.
This half passes on the old code too, so it doesn't pin the new "record before the final yield" behavior. After item_two, the list comprehension asks for one more item. That resumes the generator, and on the old code the finally records completion right then, so assert_called_once() passes either way.
In real use, StreamedResultSet stops pulling after the final chunk and never makes that extra call, which is exactly the case this PR fixes. Stopping after the item the same way makes the test fail on the old code (Called 0 times) and pass with your change. I ran both locally:
| items_done = [it async for it in resumable_done] | |
| self.assertEqual(items_done, [item_two]) | |
| mock_tracer.record_operation_completion.assert_called_once() | |
| # Take the final item and stop pulling, like StreamedResultSet does. | |
| first = await resumable_done.__anext__() | |
| self.assertEqual(first, item_two) | |
| # Completion is recorded before the final yield... | |
| mock_tracer.record_operation_completion.assert_called_once() | |
| await resumable_done.aclose() | |
| # ...and the finally doesn't record it a second time. | |
| mock_tracer.record_operation_completion.assert_called_once() |
Optionally, a last=True variant (a single PartialResultSet(last=True), with _drain_stream patched) would cover the other way stream_finished gets set.
| request_id_manager=session._database, | ||
| resource_info={"project": "p"}, | ||
| ) | ||
| self.assertEqual(list(resumable_done), [item_two]) |
There was a problem hiding this comment.
Same as the async test: list(...) makes one extra next() call, which runs the old code's finally, so this passes without the fix. The new mock-server tests already cover the sync path (3 of the 6 fail without the fix), so this one matters less, but it keeps the two test files consistent:
| self.assertEqual(list(resumable_done), [item_two]) | |
| # Take the final item and stop pulling, like StreamedResultSet does. | |
| first = next(resumable_done) | |
| self.assertEqual(first, item_two) | |
| # Completion is recorded before the final yield... | |
| mock_tracer.record_operation_completion.assert_called_once() | |
| resumable_done.close() | |
| # ...and the finally doesn't record it a second time. | |
| mock_tracer.record_operation_completion.assert_called_once() |
| # Operational latency must include the ~90ms server stream duration (>= 75ms) | ||
| # and exclude the 180ms consumer sleep (<= 165ms). | ||
| self.assertGreaterEqual(point.sum, 75) | ||
| self.assertLessEqual(point.sum, 165) |
There was a problem hiding this comment.
nit: the upper bounds here (<= 165, and <= 140 in test_operation_latency_last_true_and_last_false_consistency) leave only ~70 ms of headroom over the simulated server delay, which could flake on a slow CI runner.
Since _make_partial_result_sets doesn't set resume tokens, the whole stream is buffered before the first row is yielded. A deterministic alternative: take the first row (rows = iter(results); next(rows)) and assert that operation_latencies is already recorded before consuming the rest. The >= 75 lower bound can stay.
Description
Fixes an issue where
operation_latenciesandoperation_countclient metrics were not being emitted for streaming queries (ExecuteStreamingSql).Key Changes
_restart_on_unavailablein atry...finallyblock (sync and async insnapshot.py) to guaranteemetrics_tracer.record_operation_completion()is called when streams finish normally, exit early viabreak, or raise unhandled exceptions._StreamingResponseWrapperand_AsyncStreamingResponseWrapperinmetrics_interceptor.pyto trigger attempt metric recording when encounteringlast=Trueon partial result sets."OK"upon caught gRPC stream exceptions.last=Truechunk handling across both sync and async paths.Validation
Verified using live test runs on project
span-cloud-testing(instance=ranjithmrtest,database=gsql_enterprise_final) confirming thatoperation_latenciesandattempt_latenciesare actively emitted to Cloud Monitoring forExecuteStreamingSql.operation_latencies&attempt_latencies(p50)operation_latencies&attempt_latencies(Mean)